import os
import json
import pandas as pd
from sklearn.impute import SimpleImputer
from sklearn.preprocessing import StandardScaler, LabelEncoder
from sklearn.model_selection import train_test_split

def load_and_preprocess(json_path, test_size=0.2, random_state=0):
    # 读取 JSON 数据
    df = pd.read_json(json_path)
    # 指定 HRV 特征列（根据实际调整）
    hrv_features = [c for c in df.columns if c not in ['age', 'Gender', 'AnalysisText']]

    # 缺失值填充（均值）
    imputer = SimpleImputer(strategy='mean')
    df[hrv_features] = imputer.fit_transform(df[hrv_features])

    # 特征标准化
    scaler = StandardScaler()
    df[hrv_features] = scaler.fit_transform(df[hrv_features])

    # 构建标签
    df['age_group'] = ((df['age'] // 10) * 10).astype(int)
    le = LabelEncoder()
    df['gender_label'] = le.fit_transform(df['Gender'])

    # 特征矩阵和标签
    X = df[hrv_features].values
    y_age = df['age_group'].values
    y_gender = df['gender_label'].values

    # 数据拆分
    X_tr_age, X_te_age, y_tr_age, y_te_age = train_test_split(
        X, y_age, test_size=test_size, random_state=random_state, stratify=y_age)
    X_tr_gen, X_te_gen, y_tr_gen, y_te_gen = train_test_split(
        X, y_gender, test_size=test_size, random_state=random_state, stratify=y_gender)

    return (X_tr_age, X_te_age, y_tr_age, y_te_age,
            X_tr_gen, X_te_gen, y_tr_gen, y_te_gen,
            hrv_features)
